// web/extensions/threejs_viewer.js
import { app } from "../../scripts/app.js";

// Register the custom widget
app.registerExtension({
    name: "ThreeDViewer",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name === "ThreeDModelViewer") {
            const onExecuted = nodeType.prototype.onExecuted;
            nodeType.prototype.onExecuted = function (message) {
                onExecuted?.apply(this, arguments);

                if (message.threejs_viewer) {
                    this.addThreeJSViewer(message.threejs_viewer[0]);
                }
            };

            // Add input change listener for real-time updates
            const onPropertyChanged = nodeType.prototype.onPropertyChanged;
            nodeType.prototype.onPropertyChanged = function (name, value) {
                onPropertyChanged?.apply(this, arguments);

                // Update on ANY input change, not just camera inputs
                this.updateCameraView();
            };

            // Hook into widget callbacks for real-time updates
            const onAdded = nodeType.prototype.onAdded;
            nodeType.prototype.onAdded = function () {
                onAdded?.apply(this, arguments);
                this.setupRealtimeUpdates();

                // Add custom buttons
                this.addWidget("button", "🚀 Initialize Preview", "", () => {
                    this.initializePreview();
                });

                this.addWidget("button", "📷 Reset Camera", "", () => {
                    this.resetCameraToDefault();
                });

                this.addWidget("button", "🔄 Load/Update Model", "", () => {
                    this.loadUpdateModel();
                });
            };

            nodeType.prototype.setupRealtimeUpdates = function () {
                // Wait a bit for widgets to be created
                setTimeout(() => {
                    if (this.widgets) {
                        this.widgets.forEach(widget => {
                            if (['camera_x', 'camera_y', 'camera_z', 'lookat_x', 'lookat_y', 'lookat_z', 'camera_near', 'camera_far', 'focal_length'].includes(widget.name)) {
                                // Override the value setter
                                let _value = widget.value;
                                Object.defineProperty(widget, 'value', {
                                    get: function () { return _value; },
                                    set: function (newValue) {
                                        _value = newValue;
                                        // Trigger camera update
                                        setTimeout(() => this.node?.updateCameraView(), 0);
                                    }
                                });

                                // Hook the callback, preserving `this`
                                const originalCallback = widget.callback;
                                widget.callback = function (value) {
                                    if (originalCallback) originalCallback.call(this, value); // keep `this` = widget
                                    setTimeout(() => this.node?.updateCameraView(), 0);
                                };
                            }

                        });
                    }
                }, 100);
            };

            nodeType.prototype.resetCameraToDefault = function () {
                // Reset camera values to default
                const cameraXWidget = this.widgets.find(w => w.name === "camera_x");
                const cameraYWidget = this.widgets.find(w => w.name === "camera_y");
                const cameraZWidget = this.widgets.find(w => w.name === "camera_z");
                const lookatXWidget = this.widgets.find(w => w.name === "lookat_x");
                const lookatYWidget = this.widgets.find(w => w.name === "lookat_y");
                const lookatZWidget = this.widgets.find(w => w.name === "lookat_z");
                const cameraNearWidget = this.widgets.find(w => w.name === "camera_near");
                const cameraFarWidget = this.widgets.find(w => w.name === "camera_far");
                const focalLengthWidget = this.widgets.find(w => w.name === "focal_length");

                // Reset position and orientation
                if (cameraXWidget) cameraXWidget.value = 0.0;
                if (cameraYWidget) cameraYWidget.value = 0.0;
                if (cameraZWidget) cameraZWidget.value = 5.0;
                if (lookatXWidget) lookatXWidget.value = 0.0;
                if (lookatYWidget) lookatYWidget.value = 0.0;
                if (lookatZWidget) lookatZWidget.value = 0.0;

                // Reset camera parameters
                if (cameraNearWidget) cameraNearWidget.value = 0.1;
                if (cameraFarWidget) cameraFarWidget.value = 10.0;
                if (focalLengthWidget) focalLengthWidget.value = 50.0;

                // Update the camera view
                this.updateCameraView();

            };

            nodeType.prototype.initializePreview = async function () {
                const modelPath = this.widgets.find(w => w.name === "model_path")?.value;
                if (!modelPath) return;

                const modelInfo = await this.prepareModelFile(modelPath);
                if (modelInfo) {
                    this.addThreeJSViewer(modelInfo);
                }
            };

            nodeType.prototype.loadUpdateModel = async function () {
                const modelPath = this.widgets.find(w => w.name === "model_path")?.value;
                if (!modelPath) return;

                const modelInfo = await this.prepareModelFile(modelPath);
                if (modelInfo) {
                    this.addThreeJSViewer(modelInfo);
                }
            };

            nodeType.prototype.prepareModelFile = async function (modelPath) {
                try {
                    const response = await fetch('/threejs_prepare_model', {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({ model_path: modelPath })
                    });

                    const result = await response.json();

                    if (!result.success) {
                        console.error('Failed to prepare model:', result.error);
                        return null;
                    }

                    return {
                        path: modelPath,
                        web_url: result.web_url,
                        width: this.widgets.find(w => w.name === "canvas_width")?.value || 512,
                        height: this.widgets.find(w => w.name === "canvas_height")?.value || 512,
                        filename: result.filename,
                        camera: {
                            x: this.widgets.find(w => w.name === "camera_x")?.value || 0.0,
                            y: this.widgets.find(w => w.name === "camera_y")?.value || 0.0,
                            z: this.widgets.find(w => w.name === "camera_z")?.value || 5.0
                        },
                        lookat: {
                            x: this.widgets.find(w => w.name === "lookat_x")?.value || 0.0,
                            y: this.widgets.find(w => w.name === "lookat_y")?.value || 0.0,
                            z: this.widgets.find(w => w.name === "lookat_z")?.value || 0.0
                        }
                    };
                } catch (error) {
                    console.error('Error preparing model file:', error);
                    return null;
                }
            };

            nodeType.prototype.getModelInfoFromWidgets = function () {
                const modelPath = this.widgets.find(w => w.name === "model_path")?.value;
                if (!modelPath) return null;

                const filename = modelPath.split(/[\\/]/).pop();
                const fileHash = this.simpleHash(modelPath).substring(0, 8);

                return {
                    path: modelPath,
                    web_url: `/temp_models/${fileHash}_${filename}`,
                    width: this.widgets.find(w => w.name === "canvas_width")?.value || 512,
                    height: this.widgets.find(w => w.name === "canvas_height")?.value || 512,
                    filename: filename,
                    camera: {
                        x: this.widgets.find(w => w.name === "camera_x")?.value || 0.0,
                        y: this.widgets.find(w => w.name === "camera_y")?.value || 0.0,
                        z: this.widgets.find(w => w.name === "camera_z")?.value || 5.0
                    },
                    lookat: {
                        x: this.widgets.find(w => w.name === "lookat_x")?.value || 0.0,
                        y: this.widgets.find(w => w.name === "lookat_y")?.value || 0.0,
                        z: this.widgets.find(w => w.name === "lookat_z")?.value || 0.0
                    },
                    near: this.widgets.find(w => w.name === "camera_near")?.value || 0.1,
                    far: this.widgets.find(w => w.name === "camera_far")?.value || 1000.0,
                    focal_length: this.widgets.find(w => w.name === "focal_length")?.value || 50.0
                };
            };

            nodeType.prototype.simpleHash = function (str) {
                const encoder = new TextEncoder();
                const data = encoder.encode(str);
                let hash = 0;
                for (let i = 0; i < data.length; i++) {
                    hash = ((hash << 5) - hash + data[i]) & 0xffffffff;
                }
                return Math.abs(hash).toString(16).padStart(8, '0');
            };

            nodeType.prototype.captureAndSendImage = function () {
                const canvasWidget = this.widgets?.find(w => w.name === "threejs_canvas");
                if (!canvasWidget || !canvasWidget.element) return;

                const canvas = canvasWidget.element;
                const dataURL = canvas.toDataURL("image/png");

                fetch('/threejs_capture_image', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ dataURL: dataURL })
                }).catch(err => {
                    console.error('Failed to send captured image:', err);
                });
            };

            nodeType.prototype.captureAndSendDepth = function () {
                const canvasWidget = this.widgets?.find(w => w.name === "threejs_canvas");
                if (!canvasWidget || !canvasWidget.element) return;

                const canvas = canvasWidget.element;
                const dataURL = canvas.toDataURL("image/png");

                fetch('/threejs_capture_depth', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ dataURL: dataURL })
                }).catch(err => {
                    console.error('Failed to send captured depth:', err);
                });
            };

            nodeType.prototype.renderAndCaptureAll = function () {
                const canvasWidget = this.widgets?.find(w => w.name === "threejs_canvas");
                if (!canvasWidget || !canvasWidget.element.threeJSData) {
                    return;
                }

                const { camera, renderer, scene } = canvasWidget.element.threeJSData;

                // 1. Render normal color view
                renderer.render(scene, camera);
                this.captureAndSendImage();

                // 2. Setup depth rendering
                this.setupDepthRendering(renderer, camera, scene);

                // 3. Render and capture depth
                this.renderDepthVisualization(renderer, camera, scene);
                this.captureAndSendDepth();

                // 4. Restore normal rendering
                renderer.setRenderTarget(null);
            };

            nodeType.prototype.setupDepthRendering = function (renderer, camera, scene) {
                const canvas = renderer.domElement;

                // Create depth render target if needed
                if (!this.depthRenderTarget) {
                    this.depthRenderTarget = new THREE.WebGLRenderTarget(canvas.width, canvas.height);
                    this.depthRenderTarget.texture.minFilter = THREE.NearestFilter;
                    this.depthRenderTarget.texture.magFilter = THREE.NearestFilter;
                    this.depthRenderTarget.generateMipmaps = false;
                    this.depthRenderTarget.depthBuffer = true;
                    this.depthRenderTarget.depthTexture = new THREE.DepthTexture();
                    this.depthRenderTarget.depthTexture.type = THREE.UnsignedShortType;
                }

                // Create depth visualization material if needed
                if (!this.depthVisMaterial) {
                    this.depthVisMaterial = new THREE.ShaderMaterial({
                        uniforms: {
                            tDepth: { value: null },
                            cameraNear: { value: camera.near },
                            cameraFar: { value: camera.far }
                        },
                        vertexShader: `
                            varying vec2 vUv;
                            void main() {
                                vUv = uv;
                                gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
                            }
                        `,
                        fragmentShader: `
                            uniform sampler2D tDepth;
                            uniform float cameraNear;
                            uniform float cameraFar;
                            varying vec2 vUv;

                            float readDepth(sampler2D depthSampler, vec2 coord) {
                                float fragCoordZ = texture2D(depthSampler, coord).x;
                                float viewZ = (cameraNear * cameraFar) / ((cameraFar - cameraNear) * fragCoordZ - cameraFar);
                                return (viewZ + cameraNear) / (cameraNear - cameraFar);
                            }

                            void main() {
                                float depth = readDepth(tDepth, vUv);
                                depth = 1.0 - depth; // Invert so closer = brighter
                                gl_FragColor = vec4(vec3(depth), 1.0);
                            }
                        `
                    });
                }

                // Create fullscreen quad if needed
                if (!this.depthQuad) {
                    const geometry = new THREE.PlaneGeometry(2, 2);
                    this.depthQuad = new THREE.Mesh(geometry, this.depthVisMaterial);
                    this.depthQuadScene = new THREE.Scene();
                    this.depthQuadScene.add(this.depthQuad);
                    this.depthQuadCamera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0, 1);
                }
            };

            nodeType.prototype.renderDepthVisualization = function (renderer, camera, scene) {
                // Update camera uniforms
                this.depthVisMaterial.uniforms.cameraNear.value = camera.near;
                this.depthVisMaterial.uniforms.cameraFar.value = camera.far;

                // Render scene to depth target
                renderer.setRenderTarget(this.depthRenderTarget);
                renderer.render(scene, camera);

                // Set depth texture for visualization
                this.depthVisMaterial.uniforms.tDepth.value = this.depthRenderTarget.depthTexture;

                // Render depth visualization to main canvas
                renderer.setRenderTarget(null);
                renderer.render(this.depthQuadScene, this.depthQuadCamera);
            };

            nodeType.prototype.updateCameraView = function () {
                // Only update if we have an active Three.js scene
                const canvasWidget = this.widgets?.find(w => w.name === "threejs_canvas");
                if (!canvasWidget || !canvasWidget.element.threeJSData) {
                    return;
                }

                const { camera, renderer, scene } = canvasWidget.element.threeJSData;

                // Get current widget values with correct fallbacks
                const cameraX = this.widgets.find(w => w.name === "camera_x")?.value || 0.0;
                const cameraY = this.widgets.find(w => w.name === "camera_y")?.value || 0.0;
                const cameraZ = this.widgets.find(w => w.name === "camera_z")?.value || 5.0;
                const lookatX = this.widgets.find(w => w.name === "lookat_x")?.value || 0.0;
                const lookatY = this.widgets.find(w => w.name === "lookat_y")?.value || 0.0;
                const lookatZ = this.widgets.find(w => w.name === "lookat_z")?.value || 0.0;
                const cameraNear = this.widgets.find(w => w.name === "camera_near")?.value || 0.1;
                const cameraFar = this.widgets.find(w => w.name === "camera_far")?.value || 1000.0;
                const focalLength = this.widgets.find(w => w.name === "focal_length")?.value || 50.0;

                // Update camera position and lookAt
                camera.position.set(cameraX, cameraY, cameraZ);
                camera.lookAt(lookatX, lookatY, lookatZ);

                // Update camera near/far planes and focal length
                camera.near = cameraNear;
                camera.far = cameraFar;

                // Convert focal length to field of view and update
                const fov = 2 * Math.atan(36 / (2 * focalLength)) * (180 / Math.PI);
                camera.fov = fov;
                camera.updateProjectionMatrix();

                // RENDER AND CAPTURE BOTH COLOR AND DEPTH
                this.renderAndCaptureAll();

            };


            nodeType.prototype.addThreeJSViewer = function (modelInfo) {
                const existingViewer = this.widgets?.find(w => w.name === "threejs_canvas");
                if (existingViewer) {
                    this.removeWidget(existingViewer);
                }

                const widget = this.addDOMWidget("threejs_canvas", "canvas", document.createElement("canvas"));
                widget.element.width = modelInfo.width;
                widget.element.height = modelInfo.height;
                widget.element.style.width = modelInfo.width + "px";
                widget.element.style.height = modelInfo.height + "px";
                widget.element.style.border = "1px solid #ffffff";
                widget.element.style.backgroundColor = "#000";

                this.initThreeJS(widget.element, modelInfo);

                const minWidth = 512;
                const originalNodeHeight = 160;
                const minHeight = originalNodeHeight + modelInfo.height + 20;

                const nodeWidth = Math.max(minWidth, modelInfo.width + 40);
                const nodeHeight = Math.max(minHeight, 160 + modelInfo.height + 20);

                this.setSize([nodeWidth, nodeHeight]);
            };

            nodeType.prototype.initThreeJS = function (canvas, modelInfo) {
                if (typeof THREE === 'undefined') {
                    const script = document.createElement('script');
                    script.src = 'https://cdnjs.cloudflare.com/ajax/libs/three.js/r170/three.min.js';
                    script.onload = () => {
                        this.loadGLTFLoader(canvas, modelInfo);
                    };
                    document.head.appendChild(script);
                } else {
                    this.loadGLTFLoader(canvas, modelInfo);
                }
            };

            nodeType.prototype.loadGLTFLoader = function (canvas, modelInfo) {
                if (typeof THREE.GLTFLoader === 'undefined') {
                    // For r170, load GLTFLoader from jsdelivr
                    const script = document.createElement('script');
                    script.type = 'module';
                    script.textContent = `
                        import { GLTFLoader } from 'https://cdn.jsdelivr.net/npm/three@0.170.0/examples/jsm/loaders/GLTFLoader.js';
                        window.THREE.GLTFLoader = GLTFLoader;
                        window.dispatchEvent(new CustomEvent('GLTFLoaderReady'));
                    `;

                    window.addEventListener('GLTFLoaderReady', () => {
                        this.setupThreeJSScene(canvas, modelInfo);
                    }, { once: true });

                    document.head.appendChild(script);
                } else {
                    this.setupThreeJSScene(canvas, modelInfo);
                }
            };

            nodeType.prototype.setupThreeJSScene = function (canvas, modelInfo) {
                const scene = new THREE.Scene();
                scene.background = new THREE.Color(0x222222);

                // Convert focal length to field of view (35mm equivalent)
                const focalLength = modelInfo.focal_length || 50.0;
                const fov = 2 * Math.atan(36 / (2 * focalLength)) * (180 / Math.PI);

                const camera = new THREE.PerspectiveCamera(fov, modelInfo.width / modelInfo.height, modelInfo.near || 0.1, modelInfo.far || 1000);
                camera.position.set(modelInfo.camera.x, modelInfo.camera.y, modelInfo.camera.z);
                camera.lookAt(modelInfo.lookat.x, modelInfo.lookat.y, modelInfo.lookat.z);

                const renderer = new THREE.WebGLRenderer({ canvas: canvas, antialias: true });
                renderer.setSize(modelInfo.width, modelInfo.height);
                renderer.setPixelRatio(1);
                renderer.shadowMap.enabled = true;
                renderer.shadowMap.type = THREE.PCFSoftShadowMap;

                const ambientLight = new THREE.AmbientLight(0x404040, 10.0);
                scene.add(ambientLight);

                this.loadGLTFModel(scene, modelInfo.web_url, renderer, camera);

                const animate = () => {
                    requestAnimationFrame(animate);
                    renderer.render(scene, camera);
                };
                animate();

                // Capture initial render
                setTimeout(() => this.renderAndCaptureAll(), 100);

                canvas.threeJSData = { scene, camera, renderer };
            };

            nodeType.prototype.loadGLTFModel = function (scene, modelUrl, renderer, camera) {
                console.log(`Loading model from: ${modelUrl}`);

                const manager = new THREE.LoadingManager();
                manager.onProgress = (url, loaded, total) => {
                    console.log(`Loading progress: ${(loaded / total * 100)}%`);
                };

                manager.onLoad = () => {
                    console.log('Model loaded successfully');
                };

                manager.onError = (url) => {
                    console.error(`Failed to load: ${url}`);
                };

                const loader = new THREE.GLTFLoader(manager);

                loader.load(
                    modelUrl,
                    (gltf) => {
                        const existingModels = scene.children.filter(child =>
                            child.userData && child.userData.isLoadedModel
                        );
                        existingModels.forEach(model => scene.remove(model));

                        const model = gltf.scene;
                        model.userData.isLoadedModel = true;

                        model.traverse((child) => {
                            if (child.isMesh) {
                                child.castShadow = true;
                                child.receiveShadow = true;
                            }
                        });

                        const box = new THREE.Box3().setFromObject(model);
                        const center = box.getCenter(new THREE.Vector3());
                        const size = box.getSize(new THREE.Vector3());

                        model.position.sub(center);

                        const maxSize = Math.max(size.x, size.y, size.z);
                        if (maxSize > 3) {
                            const scale = 3 / maxSize;
                            model.scale.setScalar(scale);
                        }

                        scene.add(model);
                        console.log('Model added to scene successfully');

                        // Capture images after model is loaded
                        setTimeout(() => this.renderAndCaptureAll(), 100);
                    },
                    (progress) => {
                        if (progress.lengthComputable) {
                            const percentComplete = progress.loaded / progress.total * 100;
                            console.log(`Loading progress: ${percentComplete.toFixed(1)}%`);
                        }
                    },
                    (error) => {
                        console.error('Error loading model:', error);
                        const geometry = new THREE.BoxGeometry(1, 1, 1);
                        const material = new THREE.MeshLambertMaterial({ color: 0xff4444 });
                        const cube = new THREE.Mesh(geometry, material);
                        cube.castShadow = true;
                        cube.receiveShadow = true;
                        cube.userData.isLoadedModel = true;
                        scene.add(cube);
                        console.log('Added fallback cube due to loading error');

                        // Capture images after fallback cube is added
                        setTimeout(() => this.renderAndCaptureAll(), 100);
                    }
                );
            };
        }
    }
});
